import pandas as pd
import numpy as np

from sklearn.linear_model import Ridge
from coin.strategy.interval.profile.taekwon import greats, merge_greats
from coin.strategy.interval.util import get_df_from_total


def ridge_regress_and_predict(X_in, y_in, X_out, y_out):
  clf = Ridge(alpha=0.9, normalize=True)
  clf.fit(X_in, y_in)
  y_pred_out = clf.predict(X_out)
  return clf, y_pred_out


def cross_fit(total_df, base_product):
  ref_products = ['bitmex-xbtusd', 'okex-btcusd-quarter', 'okex-btcusd-perpetual']
  alpha_groups = merge_greats(base_product, ref_products)
  base_df = get_df_from_total(base_product, total_df)
  x_list = list()
  y = base_df['Y.vwap_ret.1m'].fillna(0)
  t = base_df['global_timestamp']
  vwap = base_df['t.vwap.1m']

  for alpha_group in alpha_groups:
    # TODO(taekwon): this is temporary. should be fixed.
    # basically, just return total_df normally
    loaded_alpha_group = alpha_group.load()
    if len(loaded_alpha_group.products) == 1:
      ref_product = loaded_alpha_group.products[0]
    else:
      ref_product = loaded_alpha_group.products[1]
    comp_df = get_df_from_total(ref_product, total_df)
    df = pd.merge(base_df,
                  comp_df,
                  how='outer',
                  on='global_timestamp',
                  suffixes=('/prod_0', '/prod_1'),
                  sort=True)

    for name, alpha_func in loaded_alpha_group.alpha_dict.items():
      x_list.append(alpha_func(df).rename(name))

  X = pd.concat(x_list, axis=1).fillna(0)
  print(X.shape)
  # print(X, y, t)

  halfidx = int(len(y) / 2)
  X_in, y_in = X[:halfidx], y[:halfidx]
  X_out, y_out = X[halfidx:], y[halfidx:]
  clf, y_pred_out = ridge_regress_and_predict(X_in, y_in, X_out, y_out)

  print(clf.coef_.shape)
  # print(X.columns)
  res_coef = dict(zip(X.columns, clf.coef_))

  res = dict()
  res['coef'] = res_coef
  res['intercept'] = clf.intercept_
  print(res)
  import json
  print(json.loads(json.dumps(res)))

  y_pred_in = clf.predict(X_in)
  print("In-Sample Score: %f" % clf.score(X_in, y_in))
  # print(np.corrcoef([y_pred_out, y_out]))
  print("Out-Of-Sample Score: %f" % clf.score(X_out, y_out))
  return np.append(y_pred_in, y_pred_out), \
    np.append(y_in, y_out), \
    t, \
    vwap, \
    f"{base_product} <- {str(ref_products)}\
    \nIS_score:{clf.score(X_in, y_in)}, OOS_score:{clf.score(X_out, y_out)}"
